#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据集管理组件
提供文件下载、重命名、存储管理等功能
"""

import os
import json
import hashlib
import time
from datetime import datetime
from urllib.parse import urlparse, urljoin
from PyQt5.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, QPushButton,
                             QLabel, QProgressBar, QListWidget, QListWidgetItem,
                             QFileDialog, QMessageBox, QLineEdit, QComboBox,
                             QGroupBox, QGridLayout, QTextEdit, QCheckBox)
from PyQt5.QtCore import Qt, QThread, pyqtSignal, QTimer, QMutex
from PyQt5.QtGui import QFont, QIcon
import requests
from pathlib import Path

class DownloadWorker(QThread):
    """下载工作线程"""
    progress_updated = pyqtSignal(int, str)  # 进度, 状态信息
    download_finished = pyqtSignal(str, bool)  # 文件路径, 是否成功
    error_occurred = pyqtSignal(str)  # 错误信息
    
    def __init__(self, url, save_path, chunk_size=8192, timeout=300):
        super().__init__()
        self.url = url
        self.save_path = save_path
        self.chunk_size = chunk_size
        self.timeout = timeout
        self.is_cancelled = False
        self.mutex = QMutex()
        
        # 设置请求头，参考百度图片爬虫
        self.headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.36',
            'Referer': 'https://image.baidu.com',
            'Accept': 'image/webp,image/apng,image/*,*/*;q=0.8',
            'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8',
            'Connection': 'keep-alive'
        }
        
    def run(self):
        """执行下载"""
        try:
            self.mutex.lock()
            if self.is_cancelled:
                return
            self.mutex.unlock()
            
            # 创建保存目录
            os.makedirs(os.path.dirname(self.save_path), exist_ok=True)
            
            # 发送请求，使用改进的请求头
            response = requests.get(self.url, headers=self.headers, stream=True, timeout=self.timeout)
            response.raise_for_status()
            
            total_size = int(response.headers.get('content-length', 0))
            downloaded_size = 0
            
            with open(self.save_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=self.chunk_size):
                    self.mutex.lock()
                    if self.is_cancelled:
                        f.close()
                        if os.path.exists(self.save_path):
                            os.remove(self.save_path)
                        return
                    self.mutex.unlock()
                    
                    if chunk:
                        f.write(chunk)
                        downloaded_size += len(chunk)
                        
                        if total_size > 0:
                            progress = int((downloaded_size / total_size) * 100)
                            self.progress_updated.emit(progress, f"下载中... {downloaded_size}/{total_size} bytes")
                        else:
                            self.progress_updated.emit(0, f"下载中... {downloaded_size} bytes")
            
            self.download_finished.emit(self.save_path, True)
            
        except Exception as e:
            self.error_occurred.emit(str(e))
            
    def cancel(self):
        """取消下载"""
        self.mutex.lock()
        self.is_cancelled = True
        self.mutex.unlock()

class DatasetManager(QWidget):
    """数据集管理组件"""
    
    # 信号定义
    download_started = pyqtSignal(str)  # 下载开始
    download_finished = pyqtSignal(str, bool)  # 文件路径, 是否成功
    download_progress = pyqtSignal(int, str)  # 进度, 状态
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self.download_workers = {}  # 存储下载工作线程
        self.download_history = []  # 下载历史
        self.setup_ui()
        self.load_config()
        
    def setup_ui(self):
        """设置UI"""
        layout = QVBoxLayout(self)
        
        # 下载设置组
        self.setup_download_settings(layout)
        
        # 下载列表组
        self.setup_download_list(layout)
        
        # 下载控制组
        self.setup_download_controls(layout)
        
    def setup_download_settings(self, parent_layout):
        """设置下载设置"""
        settings_group = QGroupBox("下载设置")
        settings_layout = QGridLayout(settings_group)
        
        # 保存目录
        settings_layout.addWidget(QLabel("保存目录:"), 0, 0)
        self.save_dir_edit = QLineEdit()
        self.save_dir_edit.setText(os.path.expanduser("~/Downloads"))
        settings_layout.addWidget(self.save_dir_edit, 0, 1)
        
        self.browse_button = QPushButton("浏览")
        self.browse_button.clicked.connect(self.browse_save_directory)
        settings_layout.addWidget(self.browse_button, 0, 2)
        
        # 最大文件大小
        settings_layout.addWidget(QLabel("最大文件大小(MB):"), 1, 0)
        self.max_size_edit = QLineEdit()
        self.max_size_edit.setText("1024")
        settings_layout.addWidget(self.max_size_edit, 1, 1)
        
        # 超时时间
        settings_layout.addWidget(QLabel("超时时间(秒):"), 2, 0)
        self.timeout_edit = QLineEdit()
        self.timeout_edit.setText("300")
        settings_layout.addWidget(self.timeout_edit, 2, 1)
        
        # 重试次数
        settings_layout.addWidget(QLabel("重试次数:"), 3, 0)
        self.retry_edit = QLineEdit()
        self.retry_edit.setText("3")
        settings_layout.addWidget(self.retry_edit, 3, 1)
        
        # 分块大小
        settings_layout.addWidget(QLabel("分块大小(KB):"), 4, 0)
        self.chunk_size_edit = QLineEdit()
        self.chunk_size_edit.setText("8")
        settings_layout.addWidget(self.chunk_size_edit, 4, 1)
        
        # 自动重命名
        self.auto_rename_check = QCheckBox("自动重命名重复文件")
        self.auto_rename_check.setChecked(True)
        settings_layout.addWidget(self.auto_rename_check, 5, 0, 1, 3)
        
        parent_layout.addWidget(settings_group)
        
    def setup_download_list(self, parent_layout):
        """设置下载列表"""
        list_group = QGroupBox("下载列表")
        list_layout = QVBoxLayout(list_group)
        
        # 下载列表
        self.download_list = QListWidget()
        list_layout.addWidget(self.download_list)
        
        # 列表控制按钮
        list_controls = QHBoxLayout()
        
        self.clear_button = QPushButton("清空列表")
        self.clear_button.clicked.connect(self.clear_download_list)
        list_controls.addWidget(self.clear_button)
        
        self.remove_button = QPushButton("移除选中")
        self.remove_button.clicked.connect(self.remove_selected_download)
        list_controls.addWidget(self.remove_button)
        
        list_controls.addStretch()
        
        list_layout.addLayout(list_controls)
        
        parent_layout.addWidget(list_group)
        
    def setup_download_controls(self, parent_layout):
        """设置下载控制"""
        controls_group = QGroupBox("下载控制")
        controls_layout = QVBoxLayout(controls_group)
        
        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setVisible(False)
        controls_layout.addWidget(self.progress_bar)
        
        # 状态标签
        self.status_label = QLabel("就绪")
        controls_layout.addWidget(self.status_label)
        
        # 控制按钮
        button_layout = QHBoxLayout()
        
        self.start_button = QPushButton("开始下载")
        self.start_button.clicked.connect(self.start_downloads)
        button_layout.addWidget(self.start_button)
        
        self.pause_button = QPushButton("暂停")
        self.pause_button.clicked.connect(self.pause_downloads)
        self.pause_button.setEnabled(False)
        button_layout.addWidget(self.pause_button)
        
        self.stop_button = QPushButton("停止")
        self.stop_button.clicked.connect(self.stop_downloads)
        button_layout.addWidget(self.stop_button)
        
        button_layout.addStretch()
        
        controls_layout.addLayout(button_layout)
        
        parent_layout.addWidget(controls_group)
        
    def browse_save_directory(self):
        """浏览保存目录"""
        directory = QFileDialog.getExistingDirectory(self, "选择保存目录")
        if directory:
            self.save_dir_edit.setText(directory)
            
    def add_download(self, url, filename=None):
        """添加下载任务"""
        if not url:
            QMessageBox.warning(self, "警告", "请输入有效的URL")
            return False
            
        # 生成文件名
        if not filename:
            filename = self.generate_filename(url)
            
        # 检查是否已存在
        for i in range(self.download_list.count()):
            item = self.download_list.item(i)
            if item.data(Qt.UserRole) == url:
                QMessageBox.information(self, "提示", "该URL已在下载列表中")
                return False
                
        # 添加到列表
        item = QListWidgetItem(f"{filename} - {url}")
        item.setData(Qt.UserRole, url)
        item.setData(Qt.UserRole + 1, filename)
        item.setData(Qt.UserRole + 2, "等待中")
        self.download_list.addItem(item)
        
        return True
        
    def generate_filename(self, url):
        """生成文件名"""
        parsed_url = urlparse(url)
        filename = os.path.basename(parsed_url.path)
        
        if not filename or '.' not in filename:
            filename = f"download_{int(time.time())}"
            
        return filename
        
    def get_save_path(self, filename):
        """获取保存路径"""
        save_dir = self.save_dir_edit.text()
        if not save_dir:
            save_dir = os.path.expanduser("~/Downloads")
            
        save_path = os.path.join(save_dir, filename)
        
        # 处理重复文件名
        if self.auto_rename_check.isChecked():
            save_path = self.get_unique_filename(save_path)
            
        return save_path
        
    def get_unique_filename(self, filepath):
        """获取唯一文件名"""
        if not os.path.exists(filepath):
            return filepath
            
        base, ext = os.path.splitext(filepath)
        counter = 1
        
        while True:
            new_path = f"{base}_{counter}{ext}"
            if not os.path.exists(new_path):
                return new_path
            counter += 1
            
    def start_downloads(self):
        """开始下载"""
        if self.download_list.count() == 0:
            QMessageBox.warning(self, "警告", "下载列表为空")
            return
            
        self.start_button.setEnabled(False)
        self.pause_button.setEnabled(True)
        self.progress_bar.setVisible(True)
        
        # 开始下载第一个任务
        self.download_next()
        
    def download_next(self):
        """下载下一个任务"""
        for i in range(self.download_list.count()):
            item = self.download_list.item(i)
            status = item.data(Qt.UserRole + 2)
            
            if status == "等待中":
                url = item.data(Qt.UserRole)
                filename = item.data(Qt.UserRole + 1)
                self.start_single_download(item, url, filename)
                break
        else:
            # 所有任务完成
            self.download_finished_all()
            
    def start_single_download(self, item, url, filename):
        """开始单个下载"""
        save_path = self.get_save_path(filename)
        
        # 更新状态
        item.setData(Qt.UserRole + 2, "下载中")
        item.setText(f"{filename} - 下载中...")
        
        # 创建下载工作线程
        worker = DownloadWorker(
            url, 
            save_path,
            chunk_size=int(self.chunk_size_edit.text()) * 1024,
            timeout=int(self.timeout_edit.text())
        )
        
        worker.progress_updated.connect(self.on_download_progress)
        worker.download_finished.connect(self.on_download_finished)
        worker.error_occurred.connect(self.on_download_error)
        
        self.download_workers[url] = worker
        worker.start()
        
        self.download_started.emit(url)
        
    def on_download_progress(self, progress, status):
        """下载进度更新"""
        self.progress_bar.setValue(progress)
        self.status_label.setText(status)
        self.download_progress.emit(progress, status)
        
    def on_download_finished(self, filepath, success):
        """下载完成"""
        # 更新列表项状态
        for i in range(self.download_list.count()):
            item = self.download_list.item(i)
            if item.data(Qt.UserRole + 2) == "下载中":
                if success:
                    item.setData(Qt.UserRole + 2, "完成")
                    item.setText(f"{os.path.basename(filepath)} - 完成")
                else:
                    item.setData(Qt.UserRole + 2, "失败")
                    item.setText(f"{os.path.basename(filepath)} - 失败")
                break
                
        # 记录下载历史
        self.download_history.append({
            'url': filepath,
            'success': success,
            'timestamp': datetime.now().isoformat()
        })
        
        self.download_finished.emit(filepath, success)
        
        # 继续下载下一个
        QTimer.singleShot(1000, self.download_next)
        
    def on_download_error(self, error_msg):
        """下载错误"""
        self.status_label.setText(f"下载错误: {error_msg}")
        QMessageBox.critical(self, "下载错误", error_msg)
        
        # 继续下载下一个
        QTimer.singleShot(1000, self.download_next)
        
    def download_finished_all(self):
        """所有下载完成"""
        self.start_button.setEnabled(True)
        self.pause_button.setEnabled(False)
        self.progress_bar.setVisible(False)
        self.status_label.setText("所有下载完成")
        
    def pause_downloads(self):
        """暂停下载"""
        for worker in self.download_workers.values():
            worker.cancel()
            
        self.pause_button.setEnabled(False)
        self.start_button.setEnabled(True)
        self.status_label.setText("下载已暂停")
        
    def stop_downloads(self):
        """停止下载"""
        for worker in self.download_workers.values():
            worker.cancel()
            
        self.download_workers.clear()
        self.pause_button.setEnabled(False)
        self.start_button.setEnabled(True)
        self.progress_bar.setVisible(False)
        self.status_label.setText("下载已停止")
        
    def clear_download_list(self):
        """清空下载列表"""
        reply = QMessageBox.question(self, "确认", "确定要清空下载列表吗？")
        if reply == QMessageBox.Yes:
            self.download_list.clear()
            
    def remove_selected_download(self):
        """移除选中的下载项"""
        current_row = self.download_list.currentRow()
        if current_row >= 0:
            item = self.download_list.takeItem(current_row)
            url = item.data(Qt.UserRole)
            
            # 如果正在下载，取消下载
            if url in self.download_workers:
                self.download_workers[url].cancel()
                del self.download_workers[url]
                
    def get_download_history(self):
        """获取下载历史"""
        return self.download_history
        
    def save_config(self):
        """保存配置"""
        config = {
            'save_directory': self.save_dir_edit.text(),
            'max_file_size': self.max_size_edit.text(),
            'timeout': self.timeout_edit.text(),
            'retry_count': self.retry_edit.text(),
            'chunk_size': self.chunk_size_edit.text(),
            'auto_rename': self.auto_rename_check.isChecked()
        }
        
        config_path = os.path.expanduser("~/.yolo_dataset_tool/config.json")
        os.makedirs(os.path.dirname(config_path), exist_ok=True)
        
        with open(config_path, 'w', encoding='utf-8') as f:
            json.dump(config, f, indent=2, ensure_ascii=False)
            
    def load_config(self):
        """加载配置"""
        config_path = os.path.expanduser("~/.yolo_dataset_tool/config.json")
        
        if os.path.exists(config_path):
            try:
                with open(config_path, 'r', encoding='utf-8') as f:
                    config = json.load(f)
                    
                self.save_dir_edit.setText(config.get('save_directory', os.path.expanduser("~/Downloads")))
                self.max_size_edit.setText(config.get('max_file_size', '1024'))
                self.timeout_edit.setText(config.get('timeout', '300'))
                self.retry_edit.setText(config.get('retry_count', '3'))
                self.chunk_size_edit.setText(config.get('chunk_size', '8'))
                self.auto_rename_check.setChecked(config.get('auto_rename', True))
                
            except Exception as e:
                print(f"加载配置失败: {e}")
                
    def get_download_statistics(self):
        """获取下载统计"""
        total = self.download_list.count()
        completed = 0
        failed = 0
        downloading = 0
        
        for i in range(total):
            item = self.download_list.item(i)
            status = item.data(Qt.UserRole + 2)
            
            if status == "完成":
                completed += 1
            elif status == "失败":
                failed += 1
            elif status == "下载中":
                downloading += 1
                
        return {
            'total': total,
            'completed': completed,
            'failed': failed,
            'downloading': downloading,
            'waiting': total - completed - failed - downloading
        }
